import json

import csv
import os
from collections import defaultdict

root = "./annotation"
save_root = "./data/json"
cFiles = os.listdir(root)
classes = ['DNoHelmetP1NoHelmetP2NoHelmet', 'DNoHelmet', 'DHelmet', 'DHelmetP1Helmet', 'DNoHelmetP0NoHelmetP1NoHelmet',
           'DNoHelmetP1NoHelmet', 'DNoHelmetP1Helmet', 'DHelmetP1NoHelmetP2Helmet', 'DHelmetP1NoHelmet',
           'DHelmetP1NoHelmetP2NoHelmet', 'DHelmetP0NoHelmetP1NoHelmetP2Helmet', 'DHelmetP0NoHelmetP1NoHelmet',
           'DNoHelmetP0NoHelmetP1NoHelmetP2NoHelmet', 'DHelmetP0NoHelmet', 'DNoHelmetP0NoHelmet',
           'DHelmetP1HelmetP2Helmet', 'DNoHelmetP1HelmetP2Helmet', 'DHelmetP0NoHelmetP1Helmet',
           'DHelmetP0NoHelmetP1NoHelmetP2NoHelmet', 'DHelmetP0HelmetP1NoHelmetP2Helmet',
           'DHelmetP0HelmetP1HelmetP2Helmet', 'DHelmetP0HelmetP1Helmet', 'DHelmetP1HelmetP2NoHelmet',
           'DNoHelmetP1NoHelmetP2NoHelmetP3NoHelmet', 'DHelmetP0Helmet', 'DHelmetP0HelmetP1NoHelmetP2NoHelmet',
           'DNoHelmetP0NoHelmetP1Helmet', 'DHelmetP0NoHelmetP1NoHelmetP2NoHelmetP3NoHelmet',
           'DNoHelmetP1NoHelmetP2Helmet', 'DHelmetP1NoHelmetP2NoHelmetP3NoHelmet',
           'DHelmetP1NoHelmetP2NoHelmetP3Helmet', 'DNoHelmetP0NoHelmetP1NoHelmetP2NoHelmetP3NoHelmet',
           'DHelmetP0NoHelmetP1HelmetP2Helmet', 'DNoHelmetP0HelmetP1NoHelmet',
           'DHelmetP0NoHelmetP1NoHelmetP2NoHelmetP3Helmet', 'DNoHelmetP0NoHelmetP1NoHelmetP2Helmet']

for i, File in enumerate(cFiles):
    jsonFile = open(save_root + '/' + File[:-4] + ".json", "w")

    annotation = defaultdict(list)
    csvFile = open(root + '/' + File, "r")
    reader = csv.reader(csvFile)
    for item in reader:
        # 忽略第一行
        if reader.line_num == 1:
            continue
        frames_id = int(item[1])
        if frames_id < 10:
            frames_id = "%02d" % frames_id
        data = [float(classes.index(item[6])), float(item[2]), float(item[3]), float(item[4]), float(item[5])]
        annotation[str(frames_id)].append(data)
    jsondata = json.dumps(annotation)
    jsonFile.write(jsondata)
    jsonFile.close()
    csvFile.close()
